from typing import Dict, List
import re
import logging

# 预定义的 Cypher 查询模板
CYPHER_TEMPLATES = {
    "find_entity": """
    MATCH (n:ProductionStage:__Entity__)
    WHERE n.id = '{entity}'
    WITH n
    OPTIONAL MATCH (n)-[r]-(m)
    OPTIONAL MATCH (n)-[r2:MENTIONS]-(d:Document)
    OPTIONAL MATCH (m)-[r3:MENTIONS]-(d2:Document)
    RETURN n, r, m, d, d2
    """,
    
    "find_all": """
    MATCH (n:ProductionStage:__Entity__)
    WITH n
    OPTIONAL MATCH (n)-[r]-(m)
    OPTIONAL MATCH (n)-[r2:MENTIONS]-(d:Document)
    RETURN n, r, m, d
    LIMIT 10
    """,
    
    "find_relation": """
    MATCH (n:ProductionStage:__Entity__)-[r]-(m)
    WHERE n.id = '{entity}'
    WITH n, r, m
    OPTIONAL MATCH (n)-[r2:MENTIONS]-(d:Document)
    OPTIONAL MATCH (m)-[r3:MENTIONS]-(d2:Document)
    RETURN n, r, m, d, d2
    """,
    
    "find_relation_between": """
    MATCH (n:ProductionStage:__Entity__), (m)
    WHERE n.id = '{entity1}' AND (m.id CONTAINS '{entity2}' OR toLower(m.id) CONTAINS toLower('{entity2}'))
    WITH n, m
    OPTIONAL MATCH (n)-[r]-(m)
    OPTIONAL MATCH (n)-[r2:MENTIONS]-(d:Document)
    OPTIONAL MATCH (m)-[r3:MENTIONS]-(d2:Document)
    RETURN n, r, m, d, d2
    """,
    
    "find_path": """
    MATCH (n:ProductionStage:__Entity__), (m)
    WHERE n.id = '{start}' AND m.id = '{end}'
    WITH n, m
    OPTIONAL MATCH path = (n)-[*..3]-(m)
    OPTIONAL MATCH (n)-[r:MENTIONS]-(d:Document)
    OPTIONAL MATCH (m)-[r2:MENTIONS]-(d2:Document)
    RETURN n, m, path, d, d2
    LIMIT 1
    """
}

def extract_entities(text: str) -> List[str]:
    """提取实体名称"""
    # 移除常见的查询词
    for word in ["查找", "搜索", "查询", "的", "相关", "关系"]:
        text = text.replace(word, " ")
    
    # 处理"与"、"和"等连接词分割
    if "与" in text:
        parts = text.split("与")
    elif "和" in text:
        parts = text.split("和")
    else:
        parts = [text]
    
    # 清理每个部分
    entities = []
    for part in parts:
        # 清理空白字符
        part = part.strip()
        if not part:
            continue
            
        # 处理特殊词
        if "温度" in part:
            if "出口物料温度" in part:
                entities.append("出口物料温度")
            elif "入口物料温度" in part:
                entities.append("入口物料温度")
            else:
                entities.append("温度")
        elif "湿度" in part:
            if "相对湿度" in part:
                entities.append("相对湿度")
            else:
                entities.append("湿度")
        elif "加料" in part:
            if "烟片加料" in part:
                entities.append("烟片加料")
            elif "叶丝加料" in part:
                entities.append("叶丝加料")
            else:
                entities.append(part)
        else:
            entities.append(part)
    
    # 清理并返回实体列表
    return [entity.strip() for entity in entities if entity.strip()]

def get_cypher_query(question: str) -> str:
    """将自然语言转换为 Cypher 查询"""
    try:
        logging.info(f"Processing question: {question}")
        question = question.lower()
        
        # 处理显示所有内容的查询
        if any(word in question for word in ["所有", "全部", "列表"]):
            logging.info("Generating query for all entities")
            return CYPHER_TEMPLATES["find_all"]
            
        # 处理关系查询
        if "关系" in question:
            entities = extract_entities(question)
            logging.info(f"Extracted entities: {entities}")
            
            if len(entities) >= 2:
                # 如果找到两个实体，使用 find_relation_between 模板
                query = CYPHER_TEMPLATES["find_relation_between"].format(
                    entity1=entities[0],
                    entity2=entities[1]
                )
            elif len(entities) == 1:
                # 如果只有一个实体，使用 find_relation 模板
                query = CYPHER_TEMPLATES["find_relation"].format(
                    entity=entities[0]
                )
            else:
                raise ValueError("未能识别出实体")
                
            logging.info(f"Generated relation query: {query}")
            return query
            
        # 处理路径查询
        if "之间" in question:
            entities = extract_entities(question)
            if len(entities) >= 2:
                query = CYPHER_TEMPLATES["find_path"].format(
                    start=entities[0],
                    end=entities[1]
                )
                logging.info(f"Generated path query: {query}")
                return query
                
        # 默认实体查询
        entities = extract_entities(question)
        if entities:
            query = CYPHER_TEMPLATES["find_entity"].format(entity=entities[0])
            logging.info(f"Generated entity query: {query}")
            return query
        else:
            raise ValueError("未能识别出实体")
                
    except Exception as e:
        logging.error(f"Error in get_cypher_query: {str(e)}")
        raise ValueError(f"无法生成 Cypher 查询: {str(e)}")